import requests
import json

# === 配置区域 ===
API_URL = "http://8.136.111.93:8008/v1/chat-messages"
API_KEY = "app-iVySNIkK8TF0MoKc7GPbCJPG"
USER_ID = "test-user-001"       # 用户唯一标识
WORKFLOW_INPUT_KEY = "query"    # 你的工作流中变量名

def ask_dify_stream(question, conv_id=None):
    """
    向 Dify Agent Chat App 发起流式请求，
    兼容基础 Chat 和 Agent Chat 的事件输出。
    """
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }
    payload = {
        "query": question,
        "inputs": {WORKFLOW_INPUT_KEY: question},
        "response_mode": "streaming",
        "user": USER_ID,
        **({"conversation_id": conv_id} if conv_id else {})
    }

    resp = requests.post(API_URL, headers=headers, json=payload, stream=True, timeout=120)
    resp.raise_for_status()

    full_answer = ""
    for line in resp.iter_lines(decode_unicode=True):
        # 只处理以 data: 开头的行
        if not line or not line.startswith("data:"):
            continue
        # 去掉 "data:" 前缀
        chunk = line[len("data:"):].strip()
        try:
            obj = json.loads(chunk)
        except json.JSONDecodeError:
            continue

        evt = obj.get("event")
        # 兼容普通 Chat 和 Agent Chat
        if evt in ("message", "agent_message"):
            text = obj.get("answer", "")
            print(text, end="", flush=True)
            full_answer += text
            # 记录 conversation_id，首次会话后可用于上下文
            conv_id = obj.get("conversation_id", conv_id)
        elif evt == "message_end":
            print()  # 流结束后换行
            break
        elif evt == "error":
            print("\n 服务端报错：", obj.get("message"))
            return conv_id, None

    return conv_id, full_answer

def main():
    print("请输入你的问题（输入 q 退出）：")
    conv_id = None
    while True:
        q = input("你问> ").strip()
        if q.lower() in ("q", "quit", "exit"):
            break
        conv_id, _ = ask_dify_stream(q, conv_id)

if __name__ == "__main__":
    main()
